{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "31fdf777",
   "metadata": {},
   "source": [
    "# Relay legalize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ce1a87e",
   "metadata": {},
   "source": [
    "Legalize pass 用于将高级算子转换为目标设备支持的形式，或替换为等效的操作序列。\n",
    "本测试文件涵盖了多种 legalize 场景，包括算子替换、无操作处理、多算子替换、多输入处理\n",
    "以及特定硬件目标下的卷积操作 legalize。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "10e1eabc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "56a09373",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pytest\n",
    "import tvm\n",
    "from tvm import te\n",
    "\n",
    "from tvm import relay\n",
    "from tvm.contrib import graph_executor\n",
    "from tvm.relay import transform, analysis\n",
    "from tvm.relay.testing.temp_op_attr import TempOpAttr\n",
    "\n",
    "def run_opt_pass(expr, passes):\n",
    "    \"\"\"\n",
    "    执行优化 passes 的通用函数。\n",
    "    \n",
    "    参数:\n",
    "    expr: 要优化的 Relay 表达式或函数\n",
    "    passes: 单个优化 pass 或优化 pass 列表\n",
    "    \n",
    "    返回值:\n",
    "    优化后的 Relay 表达式或函数\n",
    "    \"\"\"\n",
    "    passes = passes if isinstance(passes, list) else [passes]\n",
    "    mod = tvm.IRModule.from_expr(expr)\n",
    "    print(\"变换前后：\")\n",
    "    mod.show()\n",
    "    seq = tvm.transform.Sequential(passes)\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        mod = seq(mod)\n",
    "    entry = mod[\"main\"]\n",
    "    mod.show()\n",
    "    return entry if isinstance(expr, relay.Function) else entry.body"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0608a1a",
   "metadata": {},
   "source": [
    "## 算子替换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b56b884b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def before():\n",
    "    x = relay.var(\"x\", shape=(1, 64, 56, 56))\n",
    "    weight = relay.var(\"weight\", shape=(64, 64, 3, 3))\n",
    "    y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1))\n",
    "    y = relay.nn.relu(y)\n",
    "    y = relay.Function([x, weight], y)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "05f268bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def legalize_conv2d(attrs, inputs, types):\n",
    "    \"\"\"conv2d 算子的 legalize 函数，将权重乘以 2.0。\"\"\"\n",
    "    data, weight = inputs\n",
    "    weight = relay.multiply(weight, relay.const(2.0, \"float32\"))\n",
    "    return relay.nn.conv2d(data, weight, **attrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "875c3dc4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32], <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>weight, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">64</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>]);\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> multiply(<span style=\"color: #A2F; font-weight: bold\">%</span>weight, <span style=\"color: #008000\">2</span>f <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">64</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 临时设置 conv2d 的 legalize 函数\n",
    "with TempOpAttr(\"nn.conv2d\", \"FTVMLegalize\", legalize_conv2d):\n",
    "    a = before()\n",
    "    a = run_opt_pass(a, transform.Legalize())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c612b5c",
   "metadata": {},
   "source": [
    "## 测试通过返回 `None` 不做任何操作的情况"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b33865b7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32]) {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>global_max_pool2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>global_max_pool2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32]) {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>global_max_pool2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>global_max_pool2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def before():\n",
    "    \"\"\"定义 legalize 前的原始计算图。\"\"\"\n",
    "    x = relay.var(\"x\", shape=(1, 64, 56, 56))\n",
    "    y = relay.nn.global_max_pool2d(x)\n",
    "    y = relay.Function([x], y)\n",
    "    return y\n",
    "\n",
    "called = [False]\n",
    "\n",
    "def legalize_conv2d(attrs, inputs, types):\n",
    "    \"\"\"global_max_pool2d 操作的 legalize 函数，仅标记被调用，返回 None。\"\"\"\n",
    "    called[0] = True\n",
    "    return None\n",
    "\n",
    "# 临时设置 global_max_pool2d 的 legalize 函数\n",
    "with TempOpAttr(\"nn.global_max_pool2d\", \"FTVMLegalize\", legalize_conv2d):\n",
    "    a = before()\n",
    "    a = run_opt_pass(a, transform.Legalize())\n",
    "    b = run_opt_pass(before(), transform.InferType())\n",
    "\n",
    "# 验证结果与原始图相同，但 legalize 函数确实被调用了\n",
    "tvm.ir.assert_structural_equal(a, b)\n",
    "assert called[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfe10032",
   "metadata": {},
   "source": [
    "## 测试同时替换多个算子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "336280ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def before():\n",
    "    \"\"\"定义 legalize 前的原始计算图。\"\"\"\n",
    "    x = relay.var(\"x\", shape=(1, 64, 56, 56))\n",
    "    weight = relay.var(\"weight\", shape=(64, 64, 3, 3))\n",
    "    y = relay.nn.conv2d(x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1))\n",
    "    y = relay.nn.relu(y)\n",
    "    y = relay.Function([x, weight], y)\n",
    "    return y\n",
    "\n",
    "def legalize_conv2d(attrs, inputs, types):\n",
    "    \"\"\"conv2d 操作的 legalize 函数，将权重乘以 2.0。\"\"\"\n",
    "    data, weight = inputs\n",
    "    weight = relay.multiply(weight, relay.const(2.0, \"float32\"))\n",
    "    return relay.nn.conv2d(data, weight, **attrs)\n",
    "\n",
    "def legalize_relu(attrs, inputs, types):\n",
    "    \"\"\"relu 操作的 legalize 函数，在 relu 前添加一个与 0 的加法操作。\"\"\"\n",
    "    data = inputs[0]\n",
    "    add = relay.add(tvm.relay.const(0, \"float32\"), data)\n",
    "    return relay.nn.relu(add)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bfba8af8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32], <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>weight, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">64</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>]);\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> multiply(<span style=\"color: #A2F; font-weight: bold\">%</span>weight, <span style=\"color: #008000\">2</span>f <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">64</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> add(<span style=\"color: #008000\">0</span>f <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 临时设置 conv2d 和 relu 的 legalize 函数\n",
    "with TempOpAttr(\"nn.conv2d\", \"FTVMLegalize\", legalize_conv2d):\n",
    "    with TempOpAttr(\"nn.relu\", \"FTVMLegalize\", legalize_relu):\n",
    "        a = before()\n",
    "        a = run_opt_pass(a, transform.Legalize())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "994f0027",
   "metadata": {},
   "source": [
    "## 测试处理多个输入的算子的 legalize 功能"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fbc60241",
   "metadata": {},
   "outputs": [],
   "source": [
    "def before():\n",
    "    \"\"\"定义 legalize 前的原始计算图，包含多输入的 concatenate 操作。\"\"\"\n",
    "    x = relay.var(\"x\", shape=(1, 64, 56, 56))\n",
    "    y = relay.var(\"y\", shape=(1, 64, 56, 20))\n",
    "    z = relay.var(\"z\", shape=(1, 64, 56, 10))\n",
    "    func = relay.concatenate([x, y, z], axis=3)\n",
    "    func = relay.Function([x, y, z], func)\n",
    "    return func\n",
    "\n",
    "def legalize_concatenate(attrs, inputs, types):\n",
    "    \"\"\"concatenate 操作的 legalize 函数，仅验证输入类型和结构。\"\"\"\n",
    "    # 验证多输入情况的正确处理\n",
    "    assert len(inputs) == 1\n",
    "    assert isinstance(inputs[0], tvm.relay.expr.Tuple)\n",
    "    assert len(types) == 2\n",
    "    assert isinstance(types[0], tvm.relay.ty.TupleType)\n",
    "    assert isinstance(types[1], tvm.relay.ty.TensorType)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ec95a485",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32], <span style=\"color: #A2F; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">20</span>), float32], <span style=\"color: #A2F; font-weight: bold\">%</span>z: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> (<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>y, <span style=\"color: #A2F; font-weight: bold\">%</span>z);\n",
       "  concatenate(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, axis<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">3</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">20</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">20</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>z: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">86</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> (<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>y, <span style=\"color: #A2F; font-weight: bold\">%</span>z) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">56</span>), float32], Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">20</span>), float32], Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  concatenate(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, axis<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">3</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">64</span>, <span style=\"color: #008000\">56</span>, <span style=\"color: #008000\">86</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 临时设置 concatenate 的 legalize 函数\n",
    "with TempOpAttr(\"concatenate\", \"FTVMLegalize\", legalize_concatenate):\n",
    "    a = before()\n",
    "    a = run_opt_pass(a, transform.Legalize())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c79d8e2b",
   "metadata": {},
   "source": [
    "## 测试在不同 ARM CPU 目标下的卷积 2D NHWC 布局的 legalize 功能"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "180df200",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8], <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8]) {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>weight, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">4</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span>weight, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">5</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">5</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">4</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">8</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  strided_slice(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, begin<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], end<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>], strides<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>], axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8], <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8]) {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>weight, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">4</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span>weight, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">4</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">8</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  strided_slice(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, begin<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], end<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>], strides<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>], axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8], <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8]) {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>weight, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">4</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span>weight, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">5</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">5</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">4</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">8</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  strided_slice(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, begin<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], end<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>], strides<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>], axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8], <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8]) {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>weight, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">4</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span>weight, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">5</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">5</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">4</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">8</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  strided_slice(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, begin<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], end<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>], strides<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>], axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "变换前后：\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8], <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8]) {\n",
       "  nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>weight, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">4</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>weight: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span>weight, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">5</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">5</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>pad(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>int32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, pad_width<span style=\"color: #A2F; font-weight: bold\">=</span>[[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], [<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">4</span>]]) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>conv2d(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, padding<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">8</span>, kernel_size<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], data_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>, kernel_layout<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;HWIO&quot;</span>, out_dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  strided_slice(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, begin<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], end<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>], strides<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>], axes<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">None</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for target, exp_in_channels in [\n",
    "        (\n",
    "            \"llvm -device=arm_cpu -mtriple=aarch64-linux-gnu\",\n",
    "            8,\n",
    "        ),\n",
    "        (\n",
    "            \"llvm --device=arm_cpu --mtriple=aarch64-linux-gnu -mattr=+v8.2a,+dotprod\",\n",
    "            3,\n",
    "        ),\n",
    "        (\n",
    "            \"llvm --device=arm_cpu --mtriple=aarch64-linux-gnu -mattr=+v8.2a,+i8mm\",\n",
    "            8,\n",
    "        ),\n",
    "        (\n",
    "            \"llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr=+neon\",\n",
    "            8,\n",
    "        ),\n",
    "        (\n",
    "            \"llvm -device=arm_cpu -mtriple=armv8l-linux-gnu -mattr=+neon\",\n",
    "            8,\n",
    "        ),\n",
    "    ]:\n",
    "    target = tvm.target.Target(target)\n",
    "\n",
    "    dtype = \"int8\"\n",
    "    data_layout = \"NHWC\"\n",
    "    kernel_layout = \"HWIO\"\n",
    "    in_channels = 3\n",
    "    out_channels = 4\n",
    "    kernel_size = (1, 1)\n",
    "\n",
    "    x = relay.var(\"x\", shape=(1, 1, 1, in_channels), dtype=dtype)\n",
    "    weight = relay.var(\"weight\", shape=(1, 1, in_channels, out_channels), dtype=dtype)\n",
    "    out = relay.nn.conv2d(\n",
    "        x,\n",
    "        weight,\n",
    "        kernel_size=kernel_size,\n",
    "        channels=out_channels,\n",
    "        data_layout=data_layout,\n",
    "        kernel_layout=kernel_layout,\n",
    "        out_dtype=dtype,\n",
    "    )\n",
    "\n",
    "    # 在指定目标下执行 legalize pass\n",
    "    with target:\n",
    "        out = run_opt_pass(out, transform.Legalize())\n",
    "\n",
    "    # 获取 legalize 后的实际输入通道数\n",
    "    act_in_channels = out.args[0].type_args[0].shape[3]\n",
    "\n",
    "    # 验证输入通道数是否符合预期（不同的 ARM CPU 特性可能需要不同的输入通道数对齐）\n",
    "    assert act_in_channels == exp_in_channels, \"Actual input channels = \" + str(act_in_channels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd6ef9cb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
